import pathlib
import typing


class CheckpointManager:
    def __init__(self,
                 directory: pathlib.Path,
                 prefix: str = "checkpoint_",
                 suffix: str = ".pkl"):
        assert directory.is_dir()
        self._directory = directory.absolute()
        self._prefix = prefix
        self._suffix = suffix
        self._slice = slice(len(prefix), -len(suffix))

    @property
    def directory(self) -> pathlib.Path:
        return self._directory

    @property
    def prefix(self) -> str:
        return self._prefix

    @property
    def suffix(self) -> str:
        return self._suffix

    def existed(self) -> dict[int, pathlib.Path]:
        result: dict[int, pathlib.Path] = {}
        for file in self._directory.iterdir():
            if not file.is_file():
                continue
            name = file.name
            if not name.startswith(self._prefix):
                continue
            if not name.endswith(self._suffix):
                continue
            try:
                i = int(name[self._slice])
            except ValueError:
                continue
            result[i] = file
        return result

    def last_existed(self) -> None | tuple[int, pathlib.Path]:
        checkpoints = self.existed()
        if len(checkpoints) == 0:
            return None
        i = max(checkpoints)
        return i, checkpoints[i]

    _T = typing.TypeVar("_T")

    def last_existed_loaded(
            self,
            loader: typing.Callable[[pathlib.Path], _T],
            default_index: int = 0) \
            -> tuple[int, _T | None]:
        checkpoint_tuple = self.last_existed()
        if checkpoint_tuple:
            return checkpoint_tuple[0], loader(checkpoint_tuple[1])
        return default_index, None

    def new(self, i: int) -> pathlib.Path:
        return self._directory / f"{self._prefix}{i}{self._suffix}"
